"""
DatasetGenerator
"""
import numpy as np

import mindspore.dataset as ds

from .read_data import get_tensor_data, get_datalist_from_txt


class DatasetGenerator2D:
    """
    2D湍流模型数据集生成器
    用于生成包含特征数据、标签、应变率张量和雷诺应力的数据集样本
    """
    def __init__(self, data, label, sij, rs_value):
        """
        初始化数据集生成器
        
        Args:
            data: 输入特征数据
            label: 目标标签
            sij: 应变率张量数据
            rs_value: 雷诺应力数据
        """
        self.data = data
        self.label = label
        self.sij = sij
        self.rs_value = rs_value

    def __getitem__(self, index):
        """获取指定索引的样本"""
        return self.data[index], self.label[index], self.sij[index], self.rs_value[index]

    def __len__(self):
        """返回数据集样本总数"""
        return len(self.data)


def data_parallel_2d(config, data_path, rank_id, rank_size, is_train=True):
    """
    2D湍流模型的数据并行加载函数
    支持分布式训练，自动分割数据到不同计算设备
    
    Args:
        config: 配置字典
        data_path: 数据路径
        rank_id: 当前进程ID
        rank_size: 总进程数
        is_train: 是否为训练集
        
    Returns:
        mindspore.dataset: 数据集对象
    """
    # 从文本文件读取数据并转换为张量格式
    train_data_csv = get_datalist_from_txt(data_path)
    data, label, sij, re_stress = get_tensor_data(train_data_csv, config["feature_norm"],
                                                  config["label_norm"], config["data_path"])
    
    # 转换为numpy数组并调整数据类型和形状
    data = data.asnumpy().astype("float32")
    label = label.asnumpy().astype("float32").reshape(-1, 1)
    sij = sij.asnumpy().astype("float32").reshape(-1, 1)
    re_stress = re_stress.asnumpy().astype("float32").reshape(-1, 1)
    
    # 创建数据集生成器
    dataset_generator = DatasetGenerator2D(data, label, sij, re_stress)
    
    # 使用MindSpore的GeneratorDataset创建数据集，支持分布式训练
    dataset = ds.GeneratorDataset(dataset_generator, ["data", "label", "Sij", "Rs"],
                                  shuffle=True, num_shards=rank_size, shard_id=rank_id)
    
    # 批处理：训练集使用配置的批大小，测试集使用全部数据作为一批
    if is_train:
        return dataset.batch(config["batch_size"])
    return dataset.batch(len(data))


class DatasetGenerator3D:
    """
    3D湍流模型数据集生成器
    用于生成包含特征数据、标签、距离、应变率张量和雷诺应力的3D数据集样本
    """
    def __init__(self, data, label, dis, sij, rs_value):
        """
        初始化3D数据集生成器
        
        Args:
            data: 输入特征数据
            label: 目标标签
            dis: 到壁面的距离
            sij: 应变率张量数据
            rs_value: 雷诺应力数据
        """
        self.data = data
        self.label = label
        self.dis = dis
        self.sij = sij
        self.rs_value = rs_value

    def __getitem__(self, index):
        """获取指定索引的样本"""
        return self.data[index], self.label[index], self.dis[index], \
                self.sij[index], self.rs_value[index]

    def __len__(self):
        """返回数据集样本总数"""
        return len(self.data)


def data_parallel_3d(data_path, rank_id, rank_size, batch_size, is_train=True):
    """
    3D湍流模型的数据并行加载函数
    支持分布式训练，自动分割3D数据到不同计算设备
    
    Args:
        data_path: 数据路径
        rank_id: 当前进程ID
        rank_size: 总进程数
        batch_size: 批大小
        is_train: 是否为训练集
        
    Returns:
        mindspore.dataset: 数据集对象
    """
    # 从NPY文件加载3D数据并转换为float32类型
    all_data = np.load(data_path).astype(np.float32)
    
    # 解析数据：前10列为特征，第11列为标签，第12列为距离，第13列为应变率，第14列为雷诺应力
    data, label = all_data[:, 0:10], all_data[:, 10].reshape(-1, 1)
    dis, sij = all_data[:, 11].reshape(-1, 1), all_data[:, 12].reshape(-1, 1)
    re_stress = all_data[:, 13].reshape(-1, 1)
    
    # 创建3D数据集生成器
    generator = DatasetGenerator3D(data, label, dis, sij, re_stress)
    
    # 使用MindSpore的GeneratorDataset创建数据集，支持分布式训练
    columns = ["data", "label", "dis", "sij", "rs"]
    dataset = ds.GeneratorDataset(generator, columns, shuffle=True,
                                  num_shards=rank_size, shard_id=rank_id)
    
    # 批处理：训练集使用指定批大小，测试集使用全部数据作为一批
    if is_train:
        return dataset.batch(batch_size)
    return dataset.batch(len(data))